"""
首页相关的数据库操作函数
"""

from .db import get_db_connection, close_connection

async def get_statistics(user_id, start_date, end_date):
    """
    获取用户的收支统计信息
    
    Args:
        user_id: 用户ID
        start_date: 开始日期
        end_date: 结束日期
    
    Returns:
        dict: 包含收入和支出统计的字典
    """
    connection = None
    try:
        connection = await get_db_connection()
        async with connection.cursor() as cursor:
            sql = """
                SELECT 
                    COALESCE(SUM(monthly_income), 0) as income,
                    COALESCE(SUM(monthly_expense), 0) as expense
                FROM v_monthly_statistics 
                WHERE user_id = %s 
                AND month = %s
            """
            # 直接使用 start_date 的年月
            month = start_date.strftime('%Y-%m')
            await cursor.execute(sql, (user_id, month))
            result = await cursor.fetchone()
            
            return {
                "income": result['income'],
                "expense": result['expense']
            }

    except Exception as e:
        print(f"获取统计信息失败: {str(e)}")
        return None
    finally:
        await close_connection(connection)

async def get_chart_data(user_id, start_date, end_date):
    """
    获取图表数据
    
    Args:
        user_id: 用户ID
        start_date: 开始日期
        end_date: 结束日期
    
    Returns:
        dict: 包含日期、收入和支出数据的字典
    """
    connection = None
    try:
        connection = await get_db_connection()
        async with connection.cursor() as cursor:
            sql = """
                SELECT 
                    month,
                    monthly_income as income,
                    monthly_expense as expense
                FROM v_monthly_statistics 
                WHERE user_id = %s 
                AND month BETWEEN %s AND %s
                ORDER BY month DESC
                LIMIT 6
            """
            # 直接使用日期的年月
            start_month = start_date.strftime('%Y-%m')
            end_month = end_date.strftime('%Y-%m')
            await cursor.execute(sql, (user_id, start_month, end_month))
            data = await cursor.fetchall()
            
            # 处理数据格式
            dates = []
            income_data = []
            expense_data = []
            
            for row in data:
                dates.append(row['month'])
                income_data.append(float(row['income']))
                expense_data.append(float(row['expense']))
            
            return {
                "dates": dates,
                "income": income_data,
                "expense": expense_data
            }

    except Exception as e:
        print(f"获取图表数据失败: {str(e)}")
        return None
    finally:
        await close_connection(connection)
